/*
 * Kernel dynamic hooks based on ftrace
 * aurelianliu@tencent.com
 */

#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
#include <linux/module.h>
#ifndef TK5
#include <linux/kallsyms.h>
#endif
#include <linux/cpu.h>
#include <linux/delay.h>
#include <generated/utsrelease.h>
#include <asm/cacheflush.h>
#include <linux/ftrace.h>
#include <linux/string.h>
#include "version.h"
#include "hook.h"
#include "data_aware.h"

#ifdef CONFIG_FUNCTION_TRACER
#define CC_USING_FENTRY
#endif

/*
 * Patched functions
 */
#ifdef CC_USING_FENTRY
static struct verify_func verify_funcs_system[] =
{
#ifdef TK5
	{NULL, NULL, "__alloc_pages", 0, 0, (unsigned long)stat_alloc_pages_nodemask, 0},
	{NULL, NULL, "handle_mm_fault", 0, 0, (unsigned long)stat_handle_mm_fault, 0},
	{NULL, NULL, "blk_mq_dispatch_rq_list", 0, 0, (unsigned long)stat_blk_mq_dispatch_rq_list, 0},
#else
	{NULL, NULL, "__alloc_pages_nodemask", 0, 0, (unsigned long)stat_alloc_pages_nodemask, 0},
	{NULL, NULL, "blk_mq_get_driver_tag", 0, 0, (unsigned long)stat_blk_mq_get_driver_tag, 0},
#ifndef TK3
	{NULL, NULL, "do_page_fault", 0, 0, (unsigned long)stat_do_page_fault, 0},
#endif
#endif
#if defined(TK4_OLD) || defined(TK3) ||CONFIG_ARM64
	{NULL, NULL, "pick_next_task_fair", 0, 0, (unsigned long)stat_pick_next_task, 0},
#else
	{NULL, NULL, "pick_next_task", 0, 0, (unsigned long)stat_pick_next_task, 0},
#endif
	{NULL, NULL, "submit_bio", 0, 0, (unsigned long)stat_submit_bio, 0},
	{NULL, NULL, "bio_endio", 0, 0, (unsigned long)stat_bio_endio, 0},
	{NULL, NULL, "__kmalloc", 0, 0, (unsigned long)stat__kmalloc, 0},
	{NULL, NULL, "__kmalloc_node", 0, 0, (unsigned long)stat__kmalloc_node, 0},
	{NULL, NULL, "kmem_cache_alloc", 0, 0, (unsigned long)stat_kmem_cache_alloc, 0},
};
static struct verify_func verify_funcs_init[] =
{
#ifdef TK5
	{NULL, NULL, "psi_task_switch", 0, 0, (unsigned long)stat_psi_task_switch, 0},
#else
	{NULL, NULL, "finish_task_switch", 0, 0, (unsigned long)stat_finish_task_switch, 0},
#ifdef TK4_NEW
	{NULL, NULL, "sched_rqm_switch", 0, 0, (unsigned long)stat_sched_rqm_switch, 0},
#else
	{NULL, NULL, "rcu_note_context_switch", 0, 0, (unsigned long)stat_rcu_note_context_switch, 0},
#endif
#endif
};
	//{NULL, NULL, "do_syscall_64", 0, 0, (unsigned long)stat_do_syscall_64, 0},

#else

#error "error, compiler donot support fentry?";

#endif

static bool system_hooked;
int system_hook_name;
struct percpu_counter patch_num;
struct percpu_counter patch_num_system;

static struct module *get_link_module(char *modname)
{
	struct module *mod;

#ifdef TK5
	mutex_lock(module_mutex_tk5);
	mod = find_module_tk5(modname);
#else
	mutex_lock(&module_mutex);
	mod = find_module(modname);
#endif
	if (!mod) {
#ifdef TK5
		mutex_unlock(module_mutex_tk5);
#else
		mutex_unlock(&module_mutex);
#endif
		return NULL;
	}

	WARN_ON(!try_module_get(mod));
#ifdef TK5
	mutex_unlock(module_mutex_tk5);
#else
	mutex_unlock(&module_mutex);
#endif

	return mod;
}

static inline void put_link_modules(struct module *mod)
{
	if (mod)
		module_put(mod);
}

static int verify_kernel(struct verify_func verify_funcs[], int count)
{
	int i;
	char name[256];

	for (i = 0; i < count; i++) {
		if (verify_funcs[i].modname) {
			verify_funcs[i].mod = get_link_module(verify_funcs[i].modname);
			if (!verify_funcs[i].mod){
				pr_err("unable to find module '%s'\n", verify_funcs[i].modname);
				return -ENXIO;
			}
			snprintf(name, sizeof(name), "%s:%s",verify_funcs[i].modname, verify_funcs[i].name);
		} else {
			strcpy(name, verify_funcs[i].name);
		}

		/* check symbol */
#ifdef TK5
		verify_funcs[i].old_addr = kallsyms_lookup_name_tk5(name);
#else
		verify_funcs[i].old_addr = kallsyms_lookup_name(name);
#endif
		if (!verify_funcs[i].old_addr) {
			pr_err("unable to find symbol '%s'\n", name);
			return  -ENXIO;
		}
		pr_info("find symbol '%s', %lx\n", name, verify_funcs[i].old_addr);
	}
	return 0;
}

static struct verify_func *tpatch_find_match_ip(struct verify_func verify_funcs[], unsigned long ip, int count)
{
	int i;

	for ( i =0; i < count ; i++) {
		if (ip == verify_funcs[i].old_addr + verify_funcs[i].old_offset) {
			if (verify_funcs[i].modname)
				pr_info("find symbol '%s', %lx\n", verify_funcs[i].modname, verify_funcs[i].old_addr);
			return &verify_funcs[i];
		}
	}

	return NULL;
}

/* Update regs->ip to tell ftrace to return
 * to the new function.*/
#ifdef TK5
static void notrace tpatch_ftrace_handler(unsigned long ip, unsigned long parent_ip,
		      struct ftrace_ops *fops, struct ftrace_regs *regs, struct verify_func verify_funcs[], int count)
#else
static void notrace tpatch_ftrace_handler(unsigned long ip, unsigned long parent_ip,
		      struct ftrace_ops *fops, struct pt_regs *regs, struct verify_func verify_funcs[], int count)
#endif
{
	struct verify_func *func;

	preempt_disable_notrace();

	func = tpatch_find_match_ip(verify_funcs, ip, count);
	if (func)
#ifdef CONFIG_X86
#ifdef TK5
		regs->regs.ip = func->new_addr + MCOUNT_INSN_SIZE;
#else
		regs->ip = func->new_addr + MCOUNT_INSN_SIZE;
#endif
#elif defined(CONFIG_ARM64)
#ifdef TK5
		regs->regs.pc = func->new_addr + MCOUNT_INSN_SIZE;
#else
		regs->pc = func->new_addr + MCOUNT_INSN_SIZE;
#endif
#endif

	preempt_enable_notrace();
}

#ifdef TK5
static void notrace tpatch_ftrace_handler_init(unsigned long ip, unsigned long parent_ip,
		      struct ftrace_ops *fops, struct ftrace_regs *regs)
#else
static void notrace tpatch_ftrace_handler_init(unsigned long ip, unsigned long parent_ip,
		      struct ftrace_ops *fops, struct pt_regs *regs)
#endif
{
	int count;
	count = ARRAY_SIZE(verify_funcs_init);
	tpatch_ftrace_handler(ip, parent_ip, fops, regs, verify_funcs_init, count);
}
#ifdef TK5
static void notrace tpatch_ftrace_handler_system(unsigned long ip, unsigned long parent_ip,
		      struct ftrace_ops *fops, struct ftrace_regs *regs)
#else
static void notrace tpatch_ftrace_handler_system(unsigned long ip, unsigned long parent_ip,
		      struct ftrace_ops *fops, struct pt_regs *regs)
#endif
{
	int count;
	count = ARRAY_SIZE(verify_funcs_system);
	tpatch_ftrace_handler(ip, parent_ip, fops, regs, verify_funcs_system, count);
}

static struct ftrace_ops tpatch_ftrace_ops_init __read_mostly = {
	.func = tpatch_ftrace_handler_init,
	.flags = FTRACE_OPS_FL_SAVE_REGS,
};
static struct ftrace_ops tpatch_ftrace_ops_system __read_mostly = {
	.func = tpatch_ftrace_handler_system,
	.flags = FTRACE_OPS_FL_SAVE_REGS,
};

static int patch_kernel(struct ftrace_ops *tpatch_ftrace_ops, struct verify_func verify_funcs[], int count)
{
	int ret, i = 0;

	for ( i =0; i < count ; i++) {
		ret = ftrace_set_filter(tpatch_ftrace_ops, verify_funcs[i].name, strlen(verify_funcs[i].name), 0);
		if (ret < 0) {
			pr_err("can't set ftrace filter func:%s at address 0x%lx, ret(%d)\n", verify_funcs[i].name,
			       verify_funcs[i].old_addr, ret);
			goto error;
		}

		/* put module */
		put_link_modules(verify_funcs[i].mod);
	}

	ret = register_ftrace_function(tpatch_ftrace_ops);
	if (ret < 0) {
		pr_err("can't register ftrace handler\n");
		goto error;
	}

	return 0;

 error:
	for (; i < count; i++)
		put_link_modules(verify_funcs[i].mod);

	return ret;
}

int system_base_function_hook(void)
{
	int ret, count;

	if (system_hooked) {
	    pr_info("hooked already.");
	    return 0;
	}
	/* Verify patched functions */
	count = ARRAY_SIZE(verify_funcs_system);
	ret = verify_kernel(verify_funcs_system, count);
	if(ret < 0) {
		pr_err("Incorrect kernel, or function not found\n");
		return -ENODEV;
	}
#if defined(TK2) && !defined(KVM3)
	percpu_counter_init(&patch_num_system, 0);
#else
	percpu_counter_init(&patch_num_system, 0, GFP_KERNEL);
#endif

	/* Ok, try to replace target functions */
	ret = patch_kernel(&tpatch_ftrace_ops_system, verify_funcs_system, count);
	if (ret < 0)
		percpu_counter_destroy(&patch_num_system);
	else
		system_hooked = true;

	return ret;
}
int system_base_function_unhook(void)
{
	if (!system_hooked)
		return 0;

	system_hooked = false;
	/* Destroy ftrace filter */
	unregister_ftrace_function(&tpatch_ftrace_ops_system);
	synchronize_rcu();

	/* Wait all exit patched function */
	while (percpu_counter_sum(&patch_num_system))
		msleep(1);

	percpu_counter_destroy(&patch_num_system);

	return 0;
}

int __init patch_init(void)
{
	int ret, count;

	/* Verify patched functions */
	count = ARRAY_SIZE(verify_funcs_init);
	ret = verify_kernel(verify_funcs_init, count);
	if(ret < 0) {
		pr_err("Incorrect kernel, or function not found\n");
		return -ENODEV;
	}

#if defined(TK2) && !defined(KVM3)
	percpu_counter_init(&patch_num, 0);
#else
	percpu_counter_init(&patch_num, 0, GFP_KERNEL);
#endif
	/* Ok, try to replace target functions */
	ret = patch_kernel(&tpatch_ftrace_ops_init, verify_funcs_init, count);
	if (ret < 0) {
		percpu_counter_destroy(&patch_num);
		return ret;
	}

	return 0;
}

void patch_exit(void)
{
	if (system_hooked)
	   system_base_function_unhook();
	system_hooked = false;
	/* Destroy ftrace filter */
	unregister_ftrace_function(&tpatch_ftrace_ops_init);
	synchronize_rcu();

	/* Wait all exit patched function */
	while (percpu_counter_sum(&patch_num))
		msleep(1);

	percpu_counter_destroy(&patch_num);
}
